# encoding: utf-8
import torch
import traceback
from typing import List
from loguru import logger
from model.modules import UIEPredictor


class SkuExtractor(object):

    def __init__(self, task_path: str):
        self.schema = ['品牌', "型号", "通用名", "颜色", "规格"]
        self.device = "gpu" if torch.cuda.is_available() else "cpu"
        self.ie = UIEPredictor(
            model="uie-base", schema=self.schema,
            task_path=task_path,
            # task_path=f'{cfg.get("BASE_MODEL_PATH")}/model_189000',
            device=self.device, max_seq_len=256, batch_size=16
        )

    def extract(self, text: str):
        response = dict()
        response['name'] = None  # 通用名
        response['model'] = None  # 型号
        response['brand'] = None  # 品牌
        response['color'] = None  # 颜色
        response['specs'] = None  # 规格
        try:
            result = self.ie(text)
        except Exception as e:
            logger.error(traceback.format_exc())
            return response

        logger.info(result)

        if not result:
            return response
        else:
            res_dict = result[0]
            # name
            name = None
            names = res_dict.get("通用名", [])
            if names:
                name = names[0].get("text")
            model = None
            models = res_dict.get("型号", [])
            if models:
                model = models[0].get("text")
            brand = None
            brands = res_dict.get("品牌", [])
            if brands:
                brand = brands[0].get("text")
            # color
            color = None
            colors = res_dict.get("颜色", [])
            if colors:
                color = colors[0].get("text")

            # specs
            specs = None
            specses = res_dict.get("颜色", [])
            if specses:
                specs = specses[0].get("text")

            response['name'] = name
            response['model'] = model
            response['brand'] = brand
            response['color'] = color
            response['specs'] = specs

            return response

    def extract_batch(self, texts: List[str]):
        response_list = []

        result = self.ie(texts)

        for res_dict, text in zip(result, texts):
            response = dict()
            response['name'] = None  # 通用名
            response['model'] = None  # 型号
            response['brand'] = None  # 品牌
            response['color'] = None  # 颜色
            response['specs'] = None  # 规格
            # name
            name = None
            names = res_dict.get("通用名", [])
            if names:
                name = names[0].get("text")
            model = None
            models = res_dict.get("型号", [])
            if models:
                model = models[0].get("text")
            brand = None
            brands = res_dict.get("品牌", [])
            if brands:
                brand = brands[0].get("text")
            # color
            color = None
            colors = res_dict.get("颜色", [])
            if colors:
                color = colors[0].get("text")
            # specs
            specs = None
            specses = res_dict.get("颜色", [])
            if specses:
                specs = specses[0].get("text")

            response['name'] = name
            response['model'] = model
            response['brand'] = brand
            response['color'] = color
            response['specs'] = specs
            response['text'] = text

            response_list.append(response)

        return response_list
